#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <iostream>
#include <cstdlib>
#include <cstdio>
#include <cstring>
#include <chrono>
#include <memory>
using namespace std;
/**
 * 获取屏幕截图的rgb数据
 */
void getScreen(unsigned char *RGB, int w, int h)
{
    // 1.打开设备
    Display *display = XOpenDisplay(NULL);
    if (!display)
    {
        std::cerr << "Error: Unable to open X display." << std::endl;
        return;
    }
    /**
     *  这行代码是用于获取默认屏幕的屏幕编号的操作。让我解释一下它的含义：
    在X Window系统中，可以有多个屏幕（显示器）连接到计算机。
    每个屏幕都有一个唯一的编号，用于区分不同的屏幕。
    通常情况下，如果只有一个屏幕连接到计算机，它将被视为默认屏幕。
    */
    int screen_num = DefaultScreen(display);
    int screen_width = DisplayWidth(display, screen_num);   // 宽度
    int screen_height = DisplayHeight(display, screen_num); // 高度
    // 2.获取RGB32数据
    XImage *x_image = XGetImage(display, RootWindow(display, screen_num),
                                0, 0, screen_width, screen_height, AllPlanes, ZPixmap);
    if (!x_image)
    {
        std::cerr << "Error: Unable to capture screen image." << std::endl;
        return;
    }
    unsigned long red_mask = x_image->red_mask; // 这里的掩码也是4bytes，因为在XImage中ZPixmap中一个像素点也是4bytes，这里的掩码是用来萃取中对应的颜色
    unsigned long green_mask = x_image->green_mask;
    unsigned long blue_mask = x_image->blue_mask;
    int t = 0;
    for (int y = 0; y < h; ++y)
    {
        for (int x = 0; x < w; ++x)
        {
            unsigned long pixel = XGetPixel(x_image, x, y);
            uint8_t red = (pixel & red_mask) >> (8 + 8); // 这里的红色的掩码值是00ff0000，做与操作后，向右走16位，取出八位，也就是1byte的红色值

            *(RGB + t) = red;
            t++;

            uint8_t green = (pixel & green_mask) >> 8; // G绿色掩码值是0000ff00
            *(RGB + t) = green;
            t++;
            uint8_t blue = pixel & blue_mask; // 蓝色是000000ff
            *(RGB + t) = blue;
            t++;
        }
    }
}
void rgb2yuv420(unsigned char *RGB, unsigned char *yuv420, int w, int h)
{
    unsigned char *yData = yuv420;
    unsigned char *uData = yuv420 + w * h;
    unsigned char *vData = yuv420 + w * h + (int)(0.25 * w * h);
    for (int y = 0; y < h; y++)
    {
        for (int x = 0; x < w; x++)
        {
            int index = y * w + x;
            int uvIndex = (y / 2) * (w / 2) + (x / 2);
            unsigned char r = RGB[index * 3];
            unsigned char g = RGB[index * 3 + 1];
            unsigned char b = RGB[index * 3 + 2];
            int rInt = static_cast<int>(r);
            int gInt = static_cast<int>(g);
            int bInt = static_cast<int>(b);
            int yy = 0.299 * rInt + 0.587 * gInt + 0.114 * bInt;
            int u = -0.1684 * rInt - 0.3316 * gInt + 0.5 * bInt + 128;
            int v = 0.5 * rInt - 0.4187 * gInt - 0.0813 * bInt + 128;
            yy = (yy < 0) ? 0 : ((yy > 255) ? 255 : yy);
            u = (u < 0) ? 0 : ((u > 255) ? 255 : u);
            v = (v < 0) ? 0 : ((v > 255) ? 255 : v);
            yData[index] = yy;
            if (y % 2 == 0 && x % 2 == 0)
            {
                uData[uvIndex] = u;
                vData[uvIndex] = v;
            }
        }
    }
}

int main()
{
    int width = 1920, height = 1080;
    unsigned char  *RGB = new unsigned char[width * height * 3];
    unsigned char *yuv420 = new unsigned char[(int)(width * height * 1.5)];
    getScreen(RGB,width,height);
    rgb2yuv420(RGB,yuv420,width,height);
    FILE *fptr = fopen("screen.yuv420", "wb");
    fwrite(yuv420, 1, 1.5 * width * height, fptr);
    fclose(fptr);
    delete[] RGB;
    delete[] yuv420;
}